package mac5789.projeto4.algorithm;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.Map;
import java.util.Queue;
import java.util.Set;

import mac5789.projeto4.model.Action;
import mac5789.projeto4.model.Policy;
import mac5789.projeto4.model.State;
import mac5789.projeto4.model.TransitionModel;
import mac5789.projeto4.model.WorldModel;

/**
 * Implementação do Bellman Backup
 * 
 * @author lundberg
 * 
 */
public class BellmanBackup {

	private WorldModel world;

	private Set<State> activeStates;

	private Map<State, StateValue> values;

	private BellmanBackup previous;

	private BellmanBackup(WorldModel world, Map<State, StateValue> values, Set<State> activeStates, BellmanBackup previous) {
		this.world = world;
		this.values = values;
		this.activeStates = activeStates;
		this.previous = previous;
	}

	public static BellmanBackup allZeroes(WorldModel world, Set<State> activeStates) {
		Map<State, StateValue> initialValues = new HashMap<State, StateValue>();
		for (State state : world.states()) {
			initialValues.put(state, new StateValue(0.0, null));
		}
		return new BellmanBackup(world, initialValues, activeStates, null);
	}

	public BellmanBackup improve() {
		return this.improve(Collections.<State> emptySet());
	}

	public BellmanBackup improve(Collection<State> additionalActiveStates) {
		Map<State, StateValue> improved = new HashMap<State, StateValue>(this.values);
		Set<State> workingActiveStates = new HashSet<State>(this.activeStates);
		workingActiveStates.addAll(additionalActiveStates);
		for (State state : workingActiveStates) {
			StateValue improvedStateValue = evaluateImprovedState(state);
			improved.put(state, improvedStateValue);
		}
		return new BellmanBackup(world, improved, workingActiveStates, this);
	}

	private StateValue evaluateImprovedState(State state) {
		if (world.isGoal(state)) {
			return new StateValue(0.0, null);
		}
		double max = Double.NEGATIVE_INFINITY;
		Action action = null;
		for (Action candidate : world.availableActionsAt(state)) {
			double value = actionValue(state, candidate);
			if (value > max) {
				max = value;
				action = candidate;
			}
		}
		return new StateValue(world.rewardModel().reward(state) + max, action);
	}

	private double actionValue(State state, Action action) {
		double value = 0.0;
		TransitionModel transitionModel = world.transitionModelFor(state, action);
		for (State possibleState : transitionModel.possibleStates()) {
			value += transitionModel.probability(possibleState) * this.valueFor(possibleState);
		}
		value *= world.rewardModel().discountFactor();
		value -= world.rewardModel().cost(action);
		return value;
	}

	public double valueFor(State state) {
		if (!values.containsKey(state)) {
			throw new IllegalStateException("BellmanBackup não contém o valor do estado " + state);
		} else {
			return this.values.get(state).value;
		}
	}

	public Action actionFor(State state) {
		if (!values.containsKey(state)) {
			throw new IllegalStateException("BellmanBackup não contém o valor do estado " + state);
		} else {
			return this.values.get(state).action;
		}
	}

	public double getBellmanError() {
		if (previous == null) {
			return Double.POSITIVE_INFINITY;
		}
		double max = 0.0;
		for (State state : world.states()) {
			max = Math.max(max, Math.abs(this.valueFor(state) - previous.valueFor(state)));
		}
		return max;
	}

	public Policy createPolicy() {
		Map<State, Action> optimalActions = new HashMap<State, Action>();
		for (State state : this.values.keySet()) {
			optimalActions.put(state, this.actionFor(state));
		}
		return new Policy(optimalActions);
	}

	public Policy createReachablePolicy() {
		Map<State, Action> optimalActions = new HashMap<State, Action>();
		Queue<State> pending = new LinkedList<State>();
		pending.add(world.initialState());
		while (!pending.isEmpty()) {
			State state = pending.poll();
			if (!optimalActions.keySet().contains(state)) {
				Action action = this.actionFor(state);
				if (action != null) {
					pending.addAll(world.transitionModelFor(state, action).possibleStates());
					optimalActions.put(state, action);
				}
			}
		}
		return new Policy(optimalActions);
	}

	@Override
	public String toString() {
		return String.format("BellmanBackup:[values=%s,active=%s]", values, activeStates);
	}

	private static class StateValue {
		private double value;

		private Action action;

		private StateValue(double value, Action action) {
			super();
			this.value = value;
			this.action = action;
		}

		@Override
		public String toString() {
			return String.format("StateValue:[value=%s;action=%s]", value, action);
		}
	}
}
